# encoding: utf-8


"""

@author: tongzhenguo

@time: 2021/6/10 下午7:51

@desc:

 推荐系统的因果推理 | Causal Inference for Recommender Systems.pdf

"""

import causallift
from causallift import CausalLift, display
from sklearn.model_selection import train_test_split
import pandas as pd
from causallift import generate_data
import numpy as np

pd.options.display.max_rows = 8
seed = 0
data = 'simulated_observational_data'

if __name__ == "__main__":
    df = generate_data(
        N=1000,
        n_features=3,
        beta=[0, -2, 3, -5],  # Effect of [intercept and features] on outcome
        error_std=0.1,
        tau=[1, -5, -5, 10],  # Effect of [intercept and features] on treated outcome
        tau_std=0.1,
        discrete_outcome=True,
        seed=seed,
        feature_effect=0,  # Effect of beta on treated outxome
        propensity_coef=[0, -1, 1, -1],  # Effect of [intercept and features] on propensity log-odds for treatment
        index_name='index')
    print(df)
    print(pd.crosstab(df['Outcome'], df['Treatment'], margins=True))
    print(df.describe())
    print(df.corr())
    train_df, test_df = train_test_split(df, test_size=0.2, random_state=seed, stratify=df['Treatment'])
    print('\n[Estimate propensity scores for Inverse Probability Weighting.]')
    cl = CausalLift(train_df, test_df, enable_ipw=True, verbose=3)
    print('\n[Create 2 models for treatment and untreatment and estimate CATE (Conditional Average Treatment Effects)]')
    train_df, test_df = cl.estimate_cate_by_2_models()
    print('\n[Show CATE for train dataset]')
    display(train_df)
    train_df.to_csv('CATE_for_Train.csv')

    print('\n[Show CATE for test dataset]')
    display(test_df)
    test_df.to_csv('CATE_for_Test.csv')

    print('\n[Estimate the effect of recommendation based on the uplift model]')
    estimated_effect_df = cl.estimate_recommendation_impact()





